import sys
import csv

import cvxpy as cp
import numpy as np
import lib.tools as my_tools

from tqdm import tqdm
from qpsolvers import solve_qp
from lib.model import system_model
from itertools import product
from loguru import logger

# 用于多线程处理，python 3.14开始，可提供无GIL锁的多线程，是真正的并行计算
import threading
import queue

interception_system = system_model("interception_system")
logger.add("solve.log")

def solve_dataset_n3(yaml_name):
    # 加载参数
    config = my_tools.load_yaml(yaml_name)
    logger.info(f"Params {config} have been loaded.")
    eta = config['eta']
    w_u = config['w_u']
    w_omega = config['w_omega']
    
    # 状态量
    omega_z = np.linspace(-0.8, 0.8, num=int((0.8 - (-0.8)) / 0.1)+1, dtype=np.float64)
    gen_n3d = interception_system.generate_n3(bound = 0.1)
    n3d = list(gen_n3d)
    gen_n3 = interception_system.generate_n3(bound = 0.1)
    n3 = list(gen_n3)

    data = [
        ["omega_z", len(omega_z)],
        ["n3", len(n3)],
        ["n3d", len(n3d)],
        ["all combinations", len(omega_z) * len(n3) * len(n3d)]
    ]
    my_tools.print_table(data, "State seed lengths")

    # 根据状态组合数量添加带进度条的循环
    state_pdt = product(range(len(omega_z)), range(len(n3)), range(len(n3d)))
    state_len = len(omega_z) * len(n3) * len(n3d)
    state_tqdm = tqdm(state_pdt, total = state_len, ncols=110)

    processing_item = 0
    thread_nums = 7
    state_combination_list = []
    solve_result_queue = queue.Queue() # 队列用于动态存储线程返回的数据
    threads = [] # 线程列表

    with open('solved_dataset_attitude.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        # 写入表头
        writer.writerow(['D_value', 'omega_x', 'omega_y', 'omega_z', 'n3_x', 'n3_y', 'n3_z', 'n3_xd', 'n3_yd', 'n3_zd', 'desired_D_value','opt_error'])
        data_item = []
        for s1, s2, s3 in state_tqdm:

            processing_item += 1

            if processing_item % thread_nums == 0:
                state_combination_list.append([omega_z[s1], n3[s2], n3d[s3]])

                for s_cb in state_combination_list:
                    t = threading.Thread(target=n3_worker, args=(solve_result_queue, s_cb[0], s_cb[1], s_cb[2], w_omega, w_u, eta))
                    threads.append(t)
                    t.start()
                state_combination_list = []
            else:
                state_combination_list.append([omega_z[s1], n3[s2], n3d[s3]])

                if processing_item == state_len:#虽然没有凑齐目标线程数量的状态组合，但是已经是最后一个状态组合了，那么就把剩下的状态组合也创建线程处理
                    for s_cb in state_combination_list:
                        t = threading.Thread(target=n3_worker, args=(solve_result_queue, s_cb[0], s_cb[1], s_cb[2], w_omega, w_u, eta))
                        threads.append(t)
                        t.start()
                    state_combination_list = [] # 清空列表，准备存储新的thread_nums个状态组合
                else:
                    continue

        # 等待所有线程完成
        for t in threads:
            t.join()
        threads = []

        # 从队列中获取所有结果（顺序可能不一致）
        while not solve_result_queue.empty():
            data_item.append(solve_result_queue.get())
        
        my_tools.save_dataset(data_item, 'sorted_dataset_attitude.pickle')

def n3_worker(queue, omega_z, n3, n3d, w_omega, w_u, eta):
    if np.array_equal(n3, n3d):
        data_item = [0, 0, 0, omega_z, n3[0], n3[1], n3[2], n3d[0], n3d[1], n3d[2], 0, 0]
    else:
        tilde_n3x = n3[0] - n3d[0]
        tilde_n3y = n3[1] - n3d[1]
        tilde_n3z = n3[2] - n3d[2] 
        N =  np.array([[tilde_n3y*n3[2] - tilde_n3z*n3[1], tilde_n3z*n3[0] - tilde_n3x*n3[2]]])
        P = N.T @ (w_omega*N) + w_u*np.eye(2)
        n_c = -(tilde_n3x*n3[1]*omega_z - tilde_n3y*n3[0]*omega_z + eta*(tilde_n3x**2 + tilde_n3y**2))
        q = -n_c*w_omega*N.T
        lb = np.array([-2, -2])
        ub = np.array([2, 2])
        omega2 = solve_qp(P, q, lb=lb, ub=ub, solver="cvxopt")

        # 其他辅助数值计算
        D_value = interception_system.cal_D_value_n3(n_3 = n3, 
                                                    n_3d = n3d, 
                                                    omega = np.array([omega2[0], omega2[1], omega_z])
                                                    )
        error = N @ omega2.reshape(2,1) - n_c
        desired_D_value = -eta*(tilde_n3x**2 + tilde_n3y**2)
        data_item = [D_value[0,0], omega2[0], omega2[1], omega_z, n3[0], n3[1], n3[2], n3d[0], n3d[1], n3d[2], desired_D_value, error.item()]
    
    # 将结果保存在队列中
    queue.put(data_item)

COMMANDS = {
    "attitude": "Sort attitude control dataset"
}

if __name__ == "__main__":
    if len(sys.argv) == 1:
        my_tools.print_help(sys.argv[0], COMMANDS)

    if len(sys.argv) == 2:
        # param miss
        logger.error("YAML file name miss.")

    if len(sys.argv) == 3:
        if sys.argv[1] == 'attitude':
            yaml_name = sys.argv[2]
            solve_dataset_n3(yaml_name)
        else:
            my_tools.print_help(sys.argv[0], COMMANDS)